import io

import cv2
import numpy as np
from matplotlib import pyplot as plt


def rotate_roi_with_mask():
    img = cv2.imread('/Users/avi/PycharmProjects/imgperspect/WechatIMG93.jpeg')

    kernel = np.ones((1, 5), np.uint8)
    img = img[:, 400:850 + img.shape[0]]
    # 盖半行
    img[:140, 970:, :] = 0
    # 盖一行
    img[140:240, :, :] = 0
    # 再盖几行
    img[360:490, :, :] = 0
    img[620:760, :, :] = 0
    img[910: 1080, :, :] = 0
    # 转灰度
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    # 高斯模糊
    blur = cv2.GaussianBlur(gray, (25, 15), 0)
    # 直方图均衡化，防止图像过亮或过暗
    equ = cv2.equalizeHist(blur)
    # res = np.hstack((gray, equ))
    # 二值化
    ret, binary = cv2.threshold(equ, 125, 255, cv2.THRESH_BINARY)
    # 多次膨胀和腐蚀
    binary = cv2.morphologyEx(binary, cv2.MORPH_CLOSE, kernel, anchor=(2, 0), iterations=18)
    # 找边界
    contours, hierarchy = cv2.findContours(binary, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
    # cv2.drawContours(img, contours, -1, (0, 0, 255), 3)
    # cv2.imshow("blur", blur)
    # cv2.waitKey(0)
    # cv2.destroyAllWindows()
    # cv2.imshow("binary", binary)
    # cv2.waitKey(0)
    # cv2.destroyAllWindows()
    # cv2.imshow("res", res)
    # cv2.waitKey(0)
    # cv2.destroyAllWindows()

    docCnts = []

    if len(contours) > 0:
        cnts = sorted(contours, key=cv2.contourArea, reverse=True)  # 根据轮廓面积从大到小排序
        for c in cnts:
            peri = cv2.arcLength(c, True)  # 计算轮廓周长
            approx = cv2.approxPolyDP(c, 0.02 * peri, True)  # 轮廓多边形拟合
            # 如果面积特别小，就过滤掉
            if cv2.contourArea(c) < 2000:
                continue

            # 轮廓为4个点表示找到纸张
            if len(approx) == 4:
                # docCnts 中的框是由大到小排列的，保证了第一次读的肯定是一行完整麻将
                docCnts.append(approx)

    if docCnts is []:
        exit()
    # 画图
    for docCnt in docCnts:
        cv2.polylines(img, pts=[docCnt], isClosed=True, color=(0, 255, 255), thickness=3)

        for peak in docCnt:
            peak = peak[0]
            cv2.circle(img, tuple(peak), 10, (255, 0, 0))

    # cv2.imshow('img', img)
    # key = cv2.waitKey(0)

    # 图像透视变换

    img = cv2.GaussianBlur(img, (3, 3), 0)
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    # 找直线
    edges = cv2.Canny(gray, 50, 150, apertureSize=3)
    long_line_width = None
    one_dh = None
    one_dw = None
    order = 0
    for ones in docCnts:
        """
        这些角点中包括若干个完整的行（长行）由12 个麻将组成，和一个不完整行（短行），数量不确定。
        由于长行的宽度固定，则可以将长行的长边做为变换后的宽度，根据40麻将规格算出变换后的麻将高度，
        此为单张麻将变换后的高度。

        对于短行，将单张麻将变换后的高度视为已知，可以由此高度对应单张麻将的宽度（此处不准），由获取到的短行的宽度计算出有几张牌。单张麻将宽度
        乘以麻将数量，得到变换后的麻将总宽度。
        """
        order += 1
        one = np.float32(ones)[:, 0]
        xs = one[:, 0]
        ys = one[:, 1]
        # 左上角
        point_a = sorted(sorted(one, key=lambda x: x[0])[:2], key=lambda y: y[1])[0]
        # 左下角
        point_b = sorted(sorted(one, key=lambda x: x[0])[:2], key=lambda y: y[1])[1]
        # 右上角
        point_c = sorted(sorted(one, key=lambda x: x[0], reverse=True)[:2], key=lambda y: y[1])[0]
        # 右下角
        point_d = sorted(sorted(one, key=lambda x: x[0], reverse=True)[:2], key=lambda y: y[1])[1]
        # 这里的src 点的顺序要和dst 一一对应。
        src = np.float32([point_a, point_b, point_d, point_c])
        if not long_line_width:
            # 算出12 张麻将排成一排的长边有多长，只算一次，这样计算为了让捕获到的长边极值，保证更好的精度
            long_line_width = max(xs) - min(xs)
        if not one_dh:
            # 按照12 张麻将排成一排的长宽计算一张麻将的高度，只算一次
            one_dh = int(40 / 29 / 12 * long_line_width)
        if not one_dw:
            # 按照一张麻将的长宽比计算一张麻将的宽度，只算一次
            one_dw = 29 * one_dh / 40
        # 根据四个角点的值计算本次这排麻将是否是完整的，即看这排麻将的宽度是不是小于完整长度减一（0.8 是系数）
        if abs(point_d[0] - point_b[0]) < long_line_width - 0.8 * one_dw:
            # 非完整行
            # 由计算麻将张数计算宽度
            mah_num = int(abs(point_d[0] - point_b[0]) / one_dw) + 1
            transform_width = int(mah_num * one_dw)
        else:
            # 完整行
            # 计算麻将宽度
            transform_width = int(abs(point_d[0] - point_b[0]))
        dst = np.float32([[0, 0], [0, one_dh], [transform_width, one_dh], [transform_width, 0]])

        m = cv2.getPerspectiveTransform(src, dst)
        result = cv2.warpPerspective(img, m, (transform_width, one_dh))
        plt.figure()
        plt.imshow(result)
        plt.show()
        # cv2.imshow("result", result)
        # cv2.waitKey(0)
        # cv2.destroyAllWindows()
        mn = 12
        print(mn)
        for i in range(mn):
            print(f"transform_width: {transform_width}")
            calc_dw = int(transform_width / mn)
            print(f"{i} of {mn}, calc_dw: {calc_dw}, range: {i * int(calc_dw)} to {(i + 1) * int(calc_dw)}")
            save = result[:, i * int(calc_dw):(i + 1) * int(calc_dw), :]
            # cv2.imshow("save", save)
            # cv2.waitKey(0)
            # cv2.destroyAllWindows()


def rotate_roi_full_img():
    img = cv2.imread('/Users/avi/PycharmProjects/imgperspect/WechatIMG93.jpeg')

    kernel = np.ones((1, 5), np.uint8)
    img = img[:, 400:850 + img.shape[0]]
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    blur = cv2.GaussianBlur(gray, (55, 15), 0)
    equ = cv2.equalizeHist(blur)
    res = np.hstack((gray, equ))
    ret, binary = cv2.threshold(equ, 125, 255, cv2.THRESH_BINARY)
    binary = cv2.morphologyEx(binary, cv2.MORPH_CLOSE, kernel, anchor=(2, 0), iterations=18)
    contours, hierarchy = cv2.findContours(binary, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
    # cv2.drawContours(img, contours, -1, (0, 0, 255), 3)
    # cv2.imshow("blur", blur)
    # cv2.waitKey(0)
    # cv2.destroyAllWindows()
    # cv2.imshow("binary", binary)
    # cv2.waitKey(0)
    # cv2.destroyAllWindows()
    # cv2.imshow("res", res)
    # cv2.waitKey(0)
    # cv2.destroyAllWindows()

    docCnt = None

    if len(contours) > 0:
        print(len(contours))
        cnts = sorted(contours, key=cv2.contourArea, reverse=True)  # 根据轮廓面积从大到小排序
        for c in cnts:
            peri = cv2.arcLength(c, True)  # 计算轮廓周长
            approx = cv2.approxPolyDP(c, 0.02 * peri, True)  # 轮廓多边形拟合
            # 轮廓为4个点表示找到纸张
            if len(approx) == 4:
                docCnt = approx
                break
    if docCnt is None:
        exit()
    cv2.polylines(img, pts=[docCnt], isClosed=True, color=(0, 255, 255), thickness=3)

    for peak in docCnt:
        peak = peak[0]
        cv2.circle(img, tuple(peak), 10, (255, 0, 0))
    #
    # cv2.imshow('img', img)
    # key = cv2.waitKey(0)

    # rotate

    result3 = img.copy()

    img = cv2.GaussianBlur(img, (3, 3), 0)
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)

    edges = cv2.Canny(gray, 50, 150, apertureSize=3)


def parse_two_group(img, big_group, small_group) -> list:
    """
        这些角点中包括8 或9个完整的行（长行）由12 个麻将组成，和一个不完整行（短行），数量不确定。

        长行的行数由短行的高度估计。

    :param img: 操作的原始图像
    :param big_group: 8 或9个完整的行（长行）由12 个麻将组成
    :param small_group: 一个不完整行（短行）
    :return:
    """
    one = np.float32(small_group)[:, 0]
    two = np.float32(big_group)[:, 0]
    xs = one[:, 0]
    ys = one[:, 1]
    # 短行
    # 左上角
    point_a1 = sorted(sorted(one, key=lambda x: x[0])[:2], key=lambda y: y[1])[0]
    # 左下角
    point_b1 = sorted(sorted(one, key=lambda x: x[0])[:2], key=lambda y: y[1])[1]
    # 右上角
    point_c1 = sorted(sorted(one, key=lambda x: x[0], reverse=True)[:2], key=lambda y: y[1])[0]
    # 右下角
    point_d1 = sorted(sorted(one, key=lambda x: x[0], reverse=True)[:2], key=lambda y: y[1])[1]
    # 长行
    # 左上角
    point_a2 = sorted(sorted(two, key=lambda x: x[0])[:2], key=lambda y: y[1])[0]
    # 左下角
    point_b2 = sorted(sorted(two, key=lambda x: x[0])[:2], key=lambda y: y[1])[1]
    # 右上角
    point_c2 = sorted(sorted(two, key=lambda x: x[0], reverse=True)[:2], key=lambda y: y[1])[0]
    # 右下角
    point_d2 = sorted(sorted(two, key=lambda x: x[0], reverse=True)[:2], key=lambda y: y[1])[1]

    # 算短行的长度
    small_width = max((point_b1[0] - point_a1[0]), (point_d1[0] - point_c1[0]))
    # 用短行高度算出长行的麻将数
    big_width = max((point_b2[0] - point_a2[0]), (point_d2[0] - point_c2[0]))
    """
    这里是看短行长度是否小于长行的一半，如果小于一半，则认为是多张，如果多于一半，则认为是少张。这种方法要求麻将多/少张数不能超过5 张。
    """
    big_num = 8 if int(small_width / big_width) else 9
    # TODO：这里测试图片就是7 行，暂时手动设为7
    big_num = 7

    print(big_num)
    # 算长行行高
    trans_dw = point_d2[0] - point_b2[0]
    trans_dh = int(40 * big_num / 29 / 12 * trans_dw)

    res = []
    # # 长行
    # src = np.float32([point_a2, point_b2, point_d2, point_c2])
    # dst = np.float32([[0, 0], [0, trans_dh], [trans_dw, trans_dh], [trans_dw, 0]])
    # print(src)
    # print(dst)
    # m = cv2.getPerspectiveTransform(src, dst)
    # result = cv2.warpPerspective(img, m, (int(trans_dw), int(trans_dh)))
    # plt.figure()
    # plt.imshow(result)
    # plt.show()
    # fig = plt.figure()
    # seg_h = int(trans_dh / big_num)
    # seg_w = int(trans_dw / 12)
    # for hh in range(big_num):
    #     for ww in range(12):
    #         plt.subplot(big_num, 12, hh * 12 + ww + 1)
    #         seg = result[
    #               hh * seg_h:(hh + 1) * seg_h,
    #               ww * seg_w:(ww + 1) * seg_w,
    #               :]
    #         plt.imshow(seg)
    #         res.append(seg)
    #         # cv2.imwrite(str(hh*12 + ww + 1) + ".jpg", seg)
    # plt.show()
    # fig.savefig("seg_img_big.jpg")

    # 短行
    # 算长行行高: 用长行高度
    trans_dh = trans_dh / big_num
    # 按比例计算出单张宽度
    one_dw = 29 * trans_dh / 40
    # 算出短行数量: 由长行的上边长度与短行的下边长度的比例计算（想出这个真是费劲儿）
    small_num = int(np.around((point_d1[0] - point_b1[0]) / (point_c2[0] - point_a2[0]) * 12))
    trans_dw = small_num * one_dw

    src = np.float32([point_a1, point_b1, point_d1, point_c1])
    dst = np.float32([[0, 0], [0, trans_dh], [trans_dw, trans_dh], [trans_dw, 0]])
    print(src)
    print(dst)
    m = cv2.getPerspectiveTransform(src, dst)
    result = cv2.warpPerspective(img, m, (int(trans_dw), int(trans_dh)))
    plt.figure()
    plt.imshow(result)
    plt.show()
    seg_w = int(trans_dw / small_num)
    for ww in range(small_num):
        fig = plt.figure()
        f = io.BytesIO()
        # plt.subplot(1, small_num, ww + 1)
        seg = result[
              :,
              ww * seg_w:(ww + 1) * seg_w,
              :]
        plt.imshow(seg)
        plt.show()
        fig.savefig(f, format="png")
        res.append(f.getvalue())
        # cv2.imwrite(str(hh*12 + ww + 1) + ".jpg", seg)
    # fig.savefig("seg_img_small.jpg")

    return res


def parse_one_group(img, one_group) -> list:
    """
    9行12 列，硬切。

    :param img:
    :param one_group:
    :return:
    """
    src = np.float32(one_group)
    print(src)
    dw = 1500
    dh = int(40 * 9 / 29 / 12 * 1500)
    dst = np.float32([[0, 0], [0, dh], [dw, dh], [dw, 0]])
    m = cv2.getPerspectiveTransform(src, dst)
    result = cv2.warpPerspective(img, m, (dw, dh))
    # plt.figure()
    # plt.imshow(result)
    # plt.show()
    seg_h = int(result.shape[0] / 9)
    seg_w = int(result.shape[1] / 12)
    res = []
    for hh in range(9):
        for ww in range(12):
            fig = plt.figure()
            f = io.BytesIO()
            plt.subplot(9, 12, hh * 12 + ww + 1)
            seg = result[
                  hh * seg_h:(hh + 1) * seg_h,
                  ww * seg_w:(ww + 1) * seg_w,
                  :]
            plt.imshow(seg)
            plt.show()
            fig.savefig(f, format="png")
            res.append(f.getvalue())
            # cv2.imwrite(str(hh*12 + ww + 1) + ".jpg", seg)
    # fig.savefig("seg_img.jpg")
    return res


def rotate_roi_with_two_group(img_path: str) -> None:
    """
    仅处理把麻将分成两堆的情况。一堆是12 列，8 或9行；另一堆是不足12 列的一行。

    :return:
    """
    # img = cv2.imread('/Users/avi/PycharmProjects/imgperspect/WechatIMG93.jpeg')
    img = cv2.imread(img_path)
    #
    # img = img[:, 400:850 + img.shape[0]]
    # cv2.imshow("img", img)
    # cv2.waitKey(0)
    # cv2.destroyAllWindows()
    #
    # # 盖半行
    # img[:140, 970:, :] = 0
    # # 盖一行
    # img[140:240, :, :] = 0
    # # # 再盖几行
    # # img[360:490, :, :] = 0
    # # img[620:760, :, :] = 0
    # # img[910: 1080, :, :] = 0
    kernel = np.ones((1, 5), np.uint8)
    # # 转灰度
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    # 高斯模糊
    blur = cv2.GaussianBlur(gray, (25, 15), 0)
    # 直方图均衡化，防止图像过亮或过暗
    equ = cv2.equalizeHist(blur)
    # res = np.hstack((gray, equ))
    # 二值化
    ret, binary = cv2.threshold(equ, 125, 255, cv2.THRESH_BINARY)
    # 多次膨胀和腐蚀
    binary = cv2.morphologyEx(binary, cv2.MORPH_CLOSE, kernel, anchor=(2, 0), iterations=18)
    # 找边界
    contours, hierarchy = cv2.findContours(binary, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
    # cv2.drawContours(img, contours, -1, (0, 0, 255), 3)
    # cv2.imshow("blur", blur)
    # cv2.waitKey(0)
    # cv2.destroyAllWindows()
    # cv2.imshow("binary", binary)
    # cv2.waitKey(0)
    # cv2.destroyAllWindows()
    # cv2.imshow("res", res)
    # cv2.waitKey(0)
    # cv2.destroyAllWindows()

    docCnts = []

    if len(contours) > 0:
        cnts = sorted(contours, key=cv2.contourArea, reverse=True)  # 根据轮廓面积从大到小排序
        for c in cnts:
            peri = cv2.arcLength(c, True)  # 计算轮廓周长
            approx = cv2.approxPolyDP(c, 0.02 * peri, True)  # 轮廓多边形拟合
            # 如果面积特别小，就过滤掉
            if cv2.contourArea(c) < 2000:
                continue

            # 轮廓为4个点表示找到纸张
            if len(approx) == 4:
                # docCnts 中的框是由大到小排列的，保证了第一次读的肯定是一行完整麻将
                docCnts.append(approx)
            if len(docCnts) == 2:  # 只要最大的两个区域
                break

    if docCnts is []:
        return
    # 画图
    for docCnt in docCnts:
        cv2.polylines(img, pts=[docCnt], isClosed=True, color=(0, 255, 255), thickness=3)

        for peak in docCnt:
            peak = peak[0]
            cv2.circle(img, tuple(peak), 10, (255, 0, 0))
    # plt.figure()
    # plt.imshow(img)
    # plt.show()
    if len(docCnts) == 2:  # 找到了一大一小两个区域
        res = parse_two_group(img, docCnts[0], docCnts[1])
    else:  # 只有一个大区域，9 行12 列
        res = parse_one_group(img, docCnts[0])

    return res

# rotate_roi_with_mask()
# rotate_roi_full_img()
rotate_roi_with_two_group("/Users/avi/PycharmProjects/imgperspect/WechatIMG93.jpeg")

